from ultralytics.engine.model import Model
from ultralytics.nn.tasks import DetectionModel
from ultralytics.models import yolo

# 查看batchnorm层参数
def check_param(model):
    for i, (name, param) in enumerate(model.named_parameters()):
        if i == 1:
            print(f"Layer:{name},Parameter:{param}")
            break

## 模型加载测试
# 官方coco数据集.pt文件加载
pt_model = Model('./yolo11n.pt', task='detect', verbose=False)
#print(pt_model.ckpt)

# 官方.yaml配置文件加载
#pt_model._new('./yolo11n.yaml', model=DetectionModel, task='detect', verbose=True)  # 更新model为yaml配置文件
#print(pt_model.cfg)

class My_Yolo(Model):
    def __init__(self, model, task=None, verbose=False):
        super().__init__(model=model, task=task, verbose=verbose)

    @property
    def task_map(self) -> dict:
        return {'detect':{"model":DetectionModel,
                          "trainer": yolo.detect.DetectionTrainer,
                          "validator": yolo.detect.DetectionValidator,
                          "predictor": yolo.detect.DetectionPredictor,}
                }

yaml_model = My_Yolo('./yolo11n.yaml', task='detect', verbose=False)
#print(model.cfg)
'''
## 参数操作测试
model = Model('./yolo11n.pt', task='detect', verbose=True)
model.info()
check_param(model.model)
model.reset_weights()  # 重置网络层参数
check_param(model.model)

check_param(yaml_model.model)  # 检查yaml加载模型参数
yaml_model.ckpt = pt_model.ckpt  # 用官方coco的检查点伪造已训练完
yaml_model.save('./default.pt')
yaml_model.load('./yolo11n.pt')  # 导入官方coco参数，模型中断恢复
check_param(yaml_model.model)  # 检查load后的参数


if __name__=='__main__':
    train_model = My_Yolo('./yolo11n.pt', task='detect', verbose=False)
    train_model.train(data='./coco8/coco8.yaml', epochs=10)

if __name__=='__main__':
    val_model = My_Yolo('./yolo11n.pt', task='detect', verbose=False)
    val_model.val(data='./coco8/val.yaml')

if __name__=='__main__':
    pred_model = My_Yolo('./yolo11n.pt', task='detect', verbose=False)
    pred_model.predict('./coco8/images/test/000000000061.jpg', save=True)

if __name__=='__main__':
    bench_model = My_Yolo('./yolo11n.pt', task='detect', verbose=False)
    bench_model.benchmark()
'''
if __name__=='__main__':
    tune_model = My_Yolo('./yolo11n.pt', task='detect', verbose=False)
    tune_model.tune(data='./coco8/coco8.yaml', epochs=10, iterations=2)